from pymilvus import Collection, connections

import random
from pprint import pprint

connections.connect("default", host="42.194.164.248", port="19530", db_name="rensheng")
collection_name = "example_collection"
collection = Collection(collection_name)

# 定义搜索参数配置
search_params = {
    # "metric_type": "L2",  # 使用 L2 距离度量
    "metric_type": "COSINE",  # 使用 COSINE 距离度量
    "params": {"nprobe": 10}  # 设置 nprobe 参数
}

# 生成随机查询向量
query_vectors = [random.random() for _ in range(128)]

# 执行向量搜索
results = collection.search(
    data=[query_vectors],  # 查询向量列表
    anns_field="embedding",  # 向量字段名称
    param=search_params,  # 搜索参数
    limit=5,  # 返回的最近邻数量
)

pprint(results)
# 打印搜索结果
for hits in results:
    for hit in hits:
        print(f"ID: {hit.id}, Distance: {hit.distance}")
